from django.core.management.base import BaseCommand
from django.contrib.auth import get_user_model
from apps.ai_chat.models import AIProvider, AIModel, AIRole

User = get_user_model()


class Command(BaseCommand):
    help = '初始化AI对话平台的基础数据'

    def handle(self, *args, **options):
        self.stdout.write('开始初始化AI对话平台基础数据...')
        
        # 创建系统预设的AI提供商
        self.create_system_providers()
        
        # 创建系统预设的AI角色
        self.create_system_roles()
        
        self.stdout.write(
            self.style.SUCCESS('AI对话平台基础数据初始化完成！')
        )

    def create_system_providers(self):
        """创建系统预设的AI提供商"""
        self.stdout.write('创建系统预设的AI提供商...')
        
        # OpenAI提供商
        openai_provider, created = AIProvider.objects.get_or_create(
            name='OpenAI官方',
            provider_type='openai',
            defaults={
                'owner': None,  # 系统预设
                'api_endpoint': 'https://api.openai.com/v1',
                'api_key': 'your-openai-api-key-here',
                'default_params': {
                    'temperature': 0.7,
                    'max_tokens': 2000,
                    'top_p': 1.0,
                    'frequency_penalty': 0.0,
                    'presence_penalty': 0.0
                },
                'is_active': True,
                'is_public': True
            }
        )
        if created:
            self.stdout.write(f'  ✓ 创建OpenAI提供商: {openai_provider.name}')
        
        # 创建OpenAI模型
        openai_models = [
            {
                'name': 'GPT-4',
                'model_code': 'gpt-4',
                'max_tokens': 8192,
                'cost_per_1k_tokens': 0.03
            },
            {
                'name': 'GPT-4 Turbo',
                'model_code': 'gpt-4-turbo-preview',
                'max_tokens': 128000,
                'cost_per_1k_tokens': 0.01
            },
            {
                'name': 'GPT-3.5 Turbo',
                'model_code': 'gpt-3.5-turbo',
                'max_tokens': 4096,
                'cost_per_1k_tokens': 0.002
            }
        ]
        
        for model_data in openai_models:
            model, created = AIModel.objects.get_or_create(
                provider=openai_provider,
                model_code=model_data['model_code'],
                defaults=model_data
            )
            if created:
                self.stdout.write(f'    ✓ 创建模型: {model.name}')
        
        # Anthropic提供商
        anthropic_provider, created = AIProvider.objects.get_or_create(
            name='Anthropic官方',
            provider_type='anthropic',
            defaults={
                'owner': None,  # 系统预设
                'api_endpoint': 'https://api.anthropic.com/v1',
                'api_key': 'your-anthropic-api-key-here',
                'default_params': {
                    'temperature': 0.7,
                    'max_tokens': 2000
                },
                'is_active': True,
                'is_public': True
            }
        )
        if created:
            self.stdout.write(f'  ✓ 创建Anthropic提供商: {anthropic_provider.name}')
        
        # 创建Anthropic模型
        anthropic_models = [
            {
                'name': 'Claude-3 Opus',
                'model_code': 'claude-3-opus-20240229',
                'max_tokens': 200000,
                'cost_per_1k_tokens': 0.015
            },
            {
                'name': 'Claude-3 Sonnet',
                'model_code': 'claude-3-sonnet-20240229',
                'max_tokens': 200000,
                'cost_per_1k_tokens': 0.003
            },
            {
                'name': 'Claude-3 Haiku',
                'model_code': 'claude-3-haiku-20240307',
                'max_tokens': 200000,
                'cost_per_1k_tokens': 0.00025
            }
        ]
        
        for model_data in anthropic_models:
            model, created = AIModel.objects.get_or_create(
                provider=anthropic_provider,
                model_code=model_data['model_code'],
                defaults=model_data
            )
            if created:
                self.stdout.write(f'    ✓ 创建模型: {model.name}')

    def create_system_roles(self):
        """创建系统预设的AI角色"""
        self.stdout.write('创建系统预设的AI角色...')
        
        system_roles = [
            {
                'name': '通用助手',
                'system_prompt': '你是一个有用、无害、诚实的AI助手。请尽力回答用户的问题，提供准确和有帮助的信息。',
                'category': 'assistant',
                'is_public': True
            },
            {
                'name': '编程助手',
                'system_prompt': '你是一个专业的编程助手，精通多种编程语言和开发技术。请帮助用户解决编程问题，提供代码示例和最佳实践建议。',
                'category': 'programming',
                'is_public': True
            },
            {
                'name': '写作助手',
                'system_prompt': '你是一个专业的写作助手，擅长各种文体的写作。请帮助用户改进文章结构、语言表达和写作技巧。',
                'category': 'creative',
                'is_public': True
            },
            {
                'name': '数据分析师',
                'system_prompt': '你是一个专业的数据分析师，擅长数据处理、统计分析和可视化。请帮助用户分析数据，提供洞察和建议。',
                'category': 'analysis',
                'is_public': True
            },
            {
                'name': '学习导师',
                'system_prompt': '你是一个耐心的学习导师，擅长教学和知识传授。请用通俗易懂的方式解释复杂概念，帮助用户学习和理解。',
                'category': 'education',
                'is_public': True
            },
            {
                'name': '商务顾问',
                'system_prompt': '你是一个专业的商务顾问，具有丰富的商业经验。请帮助用户分析商业问题，提供战略建议和解决方案。',
                'category': 'business',
                'is_public': True
            }
        ]
        
        for role_data in system_roles:
            role, created = AIRole.objects.get_or_create(
                name=role_data['name'],
                created_by=None,  # 系统预设角色
                defaults=role_data
            )
            if created:
                self.stdout.write(f'  ✓ 创建角色: {role.name} ({role.get_category_display()})')
        
        self.stdout.write('系统预设角色创建完成！')
